// Copyright 2011 Edward Wahlen. All Rights Reserved.

package org.wahlen.pc;

import static java.lang.management.ManagementFactory.getOperatingSystemMXBean;

import com.google.common.base.Function;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;

import org.wahlen.pc.EmitFn.Collect;

import java.io.DataOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * See <a href="http://web.eecs.utk.edu/~dongarra/ccgsc2010/slides/talk28-konerding.pdf">
 * Slides</a>.
 *
 * @param <T> item type
 *
 * @author edward.wahlen@gmail (Edward Wahlen)
 */
public abstract class PCollection<T> {

  public static final int CPU_COUNT = getOperatingSystemMXBean().getAvailableProcessors();
  protected static final ExecutorService EXECUTOR = Executors.newFixedThreadPool(CPU_COUNT);

  public final PType<T> type;

  public PCollection(PType<T> type) {
    this.type = type;
  }

  protected abstract Iterator<T> shard(int i);

  /**
   * Returns a PCollection that lazily calls the given DoFn when it must. Calls the DoFn in separate
   * shards for each CPU.
   */
  @SuppressWarnings("unchecked")
  public <S> PCollection<S> parallelDo(final DoFn<T, S> f, PType<S> type) {
    return f instanceof Function ? new MapCollection<T, S>(this, (Function<T, S>) f, type)
        : new DoPCollection<T, S>(this, f, type);
  }

  /**
   * Returns a PCollection that lazily calls the given map function when it must to produce one
   * output for every input. Calls the function in separate shards for each CPU.
   */
  public <S> PCollection<S> parallelDo(final Function<T, S> f, PType<S> type) {
    return new MapCollection<T, S>(this, f, type);
  }

  /**
   * Writes the collection to a collection of files with the given string as common prefix.
   */
  public void writeTo(String prefix) throws IOException {
    try {
      EXECUTOR.invokeAll(shardWriters(prefix));
    } catch (InterruptedException e) {
      throw new AssertionError(e);
    }
  }

  private List<Callable<Void>> shardWriters(final String prefix) {
    List<Callable<Void>> result = Lists.newArrayList();
    for (int i = 0; i < CPU_COUNT; i++) {
      final int k = i;
      result.add(new Callable<Void>() {
        public Void call() throws Exception {
          DataOutputStream dos = new DataOutputStream(
              new FileOutputStream(prefix + k + "-of-" + CPU_COUNT));
          for (Iterator<T> shard = shard(k); shard.hasNext(); ) {
            type.marshaller.write(shard.next(), dos);
          }
          return null;
        }
      });
    }
    return result;
  }

  /**
   * Combines collection contents, first concurrently with separate instances created with {@code
   * Combiner.newInstance()}, and then merges results from different shards using {@code
   * Combiner.addAll()}.
   */
  public <S, C extends Combiner<T, S, C>> S combine(C combiner) {
    List<Future<C>> futures = Lists.newArrayList();
    futures.add(EXECUTOR.submit(shardWorker(combiner, 0)));
    for (int i = 1; i < CPU_COUNT; i++) {
      futures.add(EXECUTOR.submit(shardWorker(combiner.newInstance(), i)));
    }
    Iterator<Future<C>> iterator = futures.iterator();
    try {
      iterator.next().get(); // first result is the original from the method parameter.
      while (iterator.hasNext()) {
        combiner.addAll(iterator.next().get().combine());
      }
    } catch (InterruptedException e) {
      throw new AssertionError(e);
    } catch (ExecutionException e) {
      throw new AssertionError(e);
    }
    return combiner.combine();
  }

  private <C extends Combiner<T, ?, ?>> Callable<C> shardWorker(final C c, final int i) {
    return new Callable<C>() {
      @Override
      public C call() throws Exception {
        for (Iterator<T> shard = shard(i); shard.hasNext(); ) {
          c.add(shard.next());
        }
        return c;
      }
    };
  }

  /**
   * Collection based on a Function, such as {@code MapFn}. Each resulting shard simply transforms
   * items from the corresponding input shard.
   */
  private static class MapCollection<K, V> extends PCollection<V> {

    private final Function<K, V> fn;
    private final PCollection<K> collection;

    public MapCollection(PCollection<K> collection, Function<K, V> fn, PType<V> type) {
      super(type);
      this.collection = collection;
      this.fn = fn;
    }

    @Override
    protected Iterator<V> shard(int i) {
      return Iterators.transform(collection.shard(i), fn);
    }
  }

  /**
   * Collection based on a general {@code DoFn}. Each resulting shard is based on the emitted items
   * of the corresponding input shard.
   */
  private static class DoPCollection<T, S> extends PCollection<S> {

    private final PCollection<T> collection;
    private final DoFn<T, S> doFn;

    public DoPCollection(PCollection<T> collection, DoFn<T, S> doFn, PType<S> type) {
      super(type);
      this.collection = collection;
      this.doFn = doFn;
    }

    @Override
    protected Iterator<S> shard(final int i) {
      final Iterator<T> tIterator = collection.shard(i);
      doFn.start(i);
      return Iterators.concat(new Iterator<Iterator<S>>() {
        public boolean hasNext() {
          boolean result = tIterator.hasNext();
          if (!result) doFn.finish(i);
          return result;
        }

        public Iterator<S> next() {
          Collect<S> collect = new Collect<S>();
          doFn.process(tIterator.next(), collect);
          return collect.buffer.iterator();
        }

        public void remove() {
          throw new UnsupportedOperationException();
        }
      });
    }
  }
}
